#include "../lib/basic.h"
#include <cstdio>
#include <cstring>
#include <iostream>
#include <windows.h>
using namespace std;
using namespace Eigen;

const int max_thread_n = 32;

struct extr_configuration
{
	int rf_size;
	int stride_size;
	int thread_n;
	vector<string> files;
	string c_file, f_file;
};

bool parse_configuration(int argc, char *argv[], extr_configuration *config)
{
	config->thread_n = 16;
	config->stride_size = 1;
	config->files.clear();
	config->c_file = "c.bin";
	config->f_file = "f.bin";
	for (int i = 1; i < argc; ++i) {
		// thread num option
		if (strcmp(argv[i], "-t") == 0) {
			if (i + 1 >= argc) {
				printf("error with parameters: no number after '%s'!\n", argv[i]);
				return false;
			}
			++i;
			config->thread_n = s_to_i(argv[i], -1);
			if (config->thread_n < 1) {
				printf("error with parameters: invalid number after '%s'!\n", argv[i]);
				return false;
			}
			continue;
		}
		// thread num option
		if (strcmp(argv[i], "-s") == 0) {
			if (i + 1 >= argc) {
				printf("error with parameters: no number after '%s'!\n", argv[i]);
				return false;
			}
			++i;
			config->stride_size = s_to_i(argv[i], -1);
			if (config->stride_size < 1) {
				printf("error with parameters: invalid number after '%s'!\n", argv[i]);
				return false;
			}
			continue;
		}
		// centroid file option
		if (strcmp(argv[i], "-cen") == 0) {
			if (i + 1 >= argc) {
				printf("error with parameters: no <c_file> after '%s'!\n", argv[i]);
				return false;
			}
			++i;
			config->c_file = string(argv[i]);
			continue;
		}
		// feature file option
		if (strcmp(argv[i], "-fea") == 0) {
			if (i + 1 >= argc) {
				printf("error with parameters: no <f_file> after '%s'!\n", argv[i]);
				return false;
			}
			++i;
			config->f_file = string(argv[i]);
			continue;
		}
		// image file
		config->files.push_back(string(argv[i]));
	}
	return true;
}

void display_configuration(const extr_configuration &config)
{
	printf("Configuration:\n");
	printf("  rf_size=%d\n", config.rf_size);
	printf("  stride_size=%d\n", config.stride_size);
	printf("  thread_n=%d\n", config.thread_n);
	printf("  c_file='%s'\n", config.c_file.c_str());
	printf("  f_file='%s'\n", config.f_file.c_str());
	printf("  files.size()=%d\n", config.files.size());
	for (int i = 0; i < config.files.size(); ++i)
		printf("    files[%d]='%s'\n", i, config.files[i].c_str());
	printf("\n");
}

struct whitening_data
{
	bool enabled;
	RowVectorR M;
	MatrixR P;
};

void load_centroids(MatrixR &centroid, whitening_data &w_data, extr_configuration *config)
{
	FILE *f = fopen(config->c_file.c_str(), "rb");
	if (f == NULL) {
		printf("cannot open %s!\n", config->c_file.c_str());
		exit(1);
	}
	int centroid_num, patch_size;
	fread(&config->rf_size, 4, 1, f);
	fread(&centroid_num, 4, 1, f);
	fread(&patch_size, 4, 1, f);

	centroid.resize(centroid_num, patch_size);
	float *buf = new float[patch_size];
	for (int i = 0; i < centroid_num; ++i) {
		fread(buf, sizeof(float) * patch_size, 1, f);
		for (int j = 0; j < patch_size; ++j)
			centroid(i, j) = (real_num)buf[j];
	}
	delete buf;

	int wh;
	fread(&wh, 4, 1, f);
	if (wh == 1) {
		w_data.enabled = true;
		w_data.M.resize(patch_size);
		w_data.P.resize(patch_size, patch_size);
		float *buf = new float[patch_size];
		fread(buf, sizeof(float) * patch_size, 1, f);
		for (int j = 0; j < patch_size; ++j)
			w_data.M(j) = (real_num)buf[j];
		for (int i = 0; i < patch_size; ++i) {
			fread(buf, sizeof(float) * patch_size, 1, f);
			for (int j = 0; j < patch_size; ++j)
				w_data.P(i, j) = (real_num)buf[j];
		}
		delete buf;
	} else
		w_data.enabled = false;
	fclose(f);
}

void extract(const image_data &img, const MatrixR &centroid, const whitening_data &w_data, const extr_configuration &config, real_num d[])
{
	int rf_size = config.rf_size;
	int stride_size = config.stride_size;
	int centroid_num = (int)centroid.rows();
	int patch_size = (int)centroid.cols();

	int **id = new int*[img.height];
	for (int i = 0; i < img.height; ++i)
		id[i] = new int[img.width];
	int num = 0;
	for (int x = 0; x + rf_size - 1 < img.height; x += stride_size)
		for (int y = 0; y + rf_size - 1 < img.width; y += stride_size)
			id[x][y] = num++;

	printf(" - extracting patches...\n");
	real_num *p = new real_num[patch_size];
	MatrixR patch(num, patch_size);
	for (int x = 0; x + rf_size - 1 < img.height; x += stride_size)
		for (int y = 0; y + rf_size - 1 < img.width; y += stride_size) {
			extract_patch(img, x, y, rf_size, p);
			normalize_patch(p, patch_size);
			for (int j = 0; j < patch_size; ++j)
				patch(id[x][y], j) = p[j];
		}
	delete p;

	if (w_data.enabled) {
		printf(" - whitening patches...\n");
		_timer timer;
		patch.rowwise() -= w_data.M;
		//matrix_mult(patch, w_data.P, patch);
		patch *= w_data.P;
		//patch = (patch.rowwise() - w_data.M) * w_data.P;
		printf("   - mult (%d,%d) (%d,%d)  T=%s\n", (int)patch.rows(), (int)patch.cols(), (int)w_data.P.rows(), (int)w_data.P.cols(), timer.count());
	}
	
	printf(" - computing z...\n");
	_timer timer;
	VectorR p2 = patch.cwiseProduct(patch).rowwise().sum();
	VectorR c2 = centroid.cwiseProduct(centroid).rowwise().sum();
	MatrixR z; z.noalias() = patch * centroid.transpose();
	//printf("   - computed z T=%s\n",timer.count()); timer = _timer();

	z = ((-2. * z).colwise() + p2).rowwise() + c2.transpose();
	z = z.cwiseMax(0.).cwiseSqrt();
	VectorR mu = z.rowwise().sum() / (real_num)centroid_num;
	z = ((-z).colwise() + mu);
	z = z.cwiseMax(0.);

	printf("   - computed z T=%s\n",timer.count());
	
	printf(" - processing...\n");
	for (int i = 0; i < 4 * centroid_num; ++i)
		d[i] = 0.;
	for (int x = 0; x + rf_size - 1 < img.height; x += stride_size)
		for (int y = 0; y + rf_size - 1 < img.width; y += stride_size) {
			int bx = (x < (img.height - rf_size + 2) / 2 ? 0 : 1);
			int by = (y < (img.width - rf_size + 2) / 2 ? 0 : 1);
			int pos = centroid_num * (bx * 2 + by);
			//printf("x=%d y=%d (h=%d w=%d)\n", x, y, img.height, img.width);
			for (int i = 0; i < centroid_num; ++i) {
				d[pos + i] += z(id[x][y], i);
				//if (i < 10) printf("%.3lf ", z(id[x][y], i));
			}
		}

	for (int i = 0; i < img.height; ++i)
		delete id[i];
	delete id;
}

struct extract_data
{
	image_data img;
	MatrixR c;
	whitening_data w;
	extr_configuration config;
	real_num *d;
};

DWORD WINAPI extract_task(LPVOID lpParam)
{
	extract_data *data = (extract_data*)lpParam;
	extract(data->img, data->c, data->w, data->config, data->d);
	return TRUE;
}

int main(int argc, char *argv[])
{
	_timer tot_time;

	extr_configuration config;
	if (!parse_configuration(argc, argv, &config))
		return 1;

	printf("loading centroids...\n");
	MatrixR centroid;
	whitening_data w_data;
	load_centroids(centroid, w_data, &config);

	int thread_n = min(config.thread_n, max_thread_n);
	int centroid_num = (int)centroid.rows();
	int patch_size = (int)centroid.cols();
	int feature_dim = centroid_num * 4;
	display_configuration(config);

	image_data_set img;
	for (int i = 0; i < (int)config.files.size(); ++i) {
		printf("loading image %d/%d (%s) ...\n", i + 1, config.files.size(), config.files[i].c_str());
		int cnt = img.load_bin(config.files[i].c_str());
		printf("loaded %d images\n", cnt);
	}
	printf("totally loaded %d images.\n", img.size());
	if (img.size() == 0) {
		printf("no image!\n");
		return 1;
	}

	FILE *f = fopen(config.f_file.c_str(), "wb");
	if (f == NULL) {
		printf("cannot save raw features to %s!\n", config.f_file.c_str());
		exit(1);
	}
	fwrite(&feature_dim, 4, 1, f);
	
	static extract_data data[max_thread_n];
	static HANDLE threads[max_thread_n];
	for (int i = 0; i < thread_n; ++i)
		data[i].d = new real_num[feature_dim];

	printf("extracting features...\n");
	_timer timer(img.size());
	for (int cur = 0; cur < img.size(); cur += thread_n) {
		int m = min(thread_n, img.size() - cur);
		printf("extracting feature (%d~%d)/%d [remain %s]...\n", cur+1, cur+m, img.size(), timer.predict(cur));

		for (int i = 0; i < m; ++i) {
			data[i].img = img[cur + i];
			data[i].c = centroid;
			data[i].w = w_data;
			data[i].config = config;
			DWORD t_id;
			threads[i] = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)extract_task, (LPVOID)&data[i], 0, &t_id);
			if (threads[i] == NULL) {
				printf(" - unable to create thread!!! error=%d\n", GetLastError());
				exit(1);
			}
		}
		WaitForMultipleObjects(m, threads, TRUE, INFINITE);
		for (int i = 0; i < m; ++i)
			CloseHandle(threads[i]);

		for (int k = 0; k < m; ++k) {
			fwrite(&data[k].img.label, 4, 1, f);
			for (int j = 0; j < feature_dim; ++j) {
				float x = (float)data[k].d[j];
				fwrite(&x, sizeof(float), 1, f);
			}
			//printf("%d", data[k].img.label);
			//for (int j = 0; j < 5; ++j) printf(" %.3lf", data[k].d[j]); printf("\n");
		}
	}
	fclose(f);

	printf("done. [time=%s]\n", tot_time.count());
	freopen("__extract_time.txt", "w", stdout);
	printf("done. [time=%s]\n", tot_time.count());

	return 0;
}

